import socket
import time
import yaml
import cv2
import json
import numpy as np

###############################################################################
#                                  依赖库
###############################################################################
# pip install pyyaml

track_line_flag = False  # 是否显示轨迹线

###############################################################################
#                                socket 配置
###############################################################################
listen_ip = "192.168.170.41"
listen_port = 5050
socket_fd = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
socket_fd.settimeout(5)
socket_fd.bind((listen_ip, listen_port))

###############################################################################
#                               加载 PGM 图片
###############################################################################
image_path = 'out.pgm'
gray = cv2.imread(image_path, cv2.IMREAD_GRAYSCALE)
bak = cv2.merge([gray, gray, gray])
height, width = gray.shape

###############################################################################
#                                函数定义区
###############################################################################
# 将像素坐标转化为带有方向的三角形的3个顶点
def triangle_coord(x, y, a):
    angle = 210                             # 角度值, 0度正右, 顺时针旋转
    line = 5                                # 三角形大小尺寸
    t1 = a
    x1 = x + 2 * line * np.cos(t1)
    y1 = y + 2 * line * np.sin(t1)
    t2 = a + np.pi / 3 * 2
    x2 = x + line * np.cos(t2)
    y2 = y + line * np.sin(t2)
    t3 = a + np.pi / 3 * 4
    x3 = x + line * np.cos(t3)
    y3 = y + line * np.sin(t3)
    points = np.array([[x1, y1], [x2, y2], [x3, y3]], np.int32)
    return points

###############################################################################
#                              解析 YAML 配置
###############################################################################
with open('out.yaml', 'r') as file:
    yamldata = yaml.safe_load(file)
    resolution = yamldata.get('resolution')
    origin = yamldata.get('origin')

# 函数定义，3d 坐标转 2d 坐标，输出 2d 像素坐标
def dot_pix_coord(pcdx, pcdy):
    x = int((pcdx - origin[0]) / resolution)
    y = height - int((pcdy - origin[1]) / resolution)
    return x, y

###############################################################################
#                              轨迹线相关
###############################################################################

if track_line_flag:
    ###############################################################################
    #                            在背景图中添加巡检轨迹线
    ###############################################################################
    px1, py1 = dot_pix_coord(-0.132258, -0.00956725 )
    px2, py2 = dot_pix_coord( 3.7373  , -0.963677   )
    px3, py3 = dot_pix_coord( 3.44093 , -5.4706     )
    px4, py4 = dot_pix_coord(-2.54368 , -5.33527    )
    print("node1(", px1, ",", py1, ")")
    print("node2(", px2, ",", py2, ")")
    print("node3(", px3, ",", py3, ")")
    print("node4(", px4, ",", py4, ")")

    track_thickness = 1
    track_color = (255, 0, 0)
    cv2.arrowedLine(bak, (px1, py1), (px2, py2), track_color, track_thickness, tipLength=0.1)
    cv2.arrowedLine(bak, (px2, py2), (px3, py3), track_color, track_thickness, tipLength=0.1)
    cv2.arrowedLine(bak, (px3, py3), (px4, py4), track_color, track_thickness, tipLength=0.1)
    cv2.arrowedLine(bak, (px4, py4), (px1, py1), track_color, track_thickness, tipLength=0.1)

###############################################################################
#                          绘制第一帧，机器狗默认在原点
###############################################################################
points = triangle_coord(0, 0, 0)
img = bak.copy()
cv2.fillPoly(img, [points], color=(0, 0, 255))
cv2.imshow('PGM Image', img)

while True:
    try:
        udpdata, server = socket_fd.recvfrom(1024)
    except socket.timeout:
        print("no data received, timeout occurred.")
        continue

    val = json.loads(udpdata.decode("utf-8"))
    gx = val['x']
    gy = val['y']
    ga = val['a']

    # print(gx, gy, ga)

    img = bak.copy()
    points = triangle_coord(gx, gy, ga)
    cv2.fillPoly(img, [points], color=(0, 0, 255))
    cv2.imshow('PGM Image', img)

    if cv2.waitKey(1) & 0xFF == ord('q'):
        break  

cv2.destroyAllWindows()